import os
from PIL import Image
import numpy as np
import time
import logging
import u2net

def naive_cutout(img, mask):
    empty = Image.new("RGBA", (img.size), (255,255,255,255))

    img.save("./output_img/" + "111-1.png")
    mask.save("./output_img/" + "111-mask.png")
    empty.save("./output_img/" + "111-empty.png")

    cutout = Image.composite(img, empty, mask.resize(img.size, Image.LANCZOS))
    return cutout


def backgroundRemoval(image_path):
    # Get original dataset
    img = Image.open(image_path)

    # Process Image
    mask = u2net.run(np.array(img)).convert("L")

    # masking
    img_bg_removed = naive_cutout(img, mask)
    img_bg_removed.save("./output_img/" + "111.png")

if __name__ == '__main__':
    backgroundRemoval("./input_img/test2.png")
